from flask import Flask, request, jsonify
import io
import base64
from PIL import Image
from io import BytesIO
from cartoon_generator import CartoonGenerator
from torchvision import transforms
from generate_main import generate_image
import time
from logging_config import setup_logger

# 设置日志记录器
logger = setup_logger(__name__)

app = Flask(__name__)

# 将device和 CartoonGenerator 初始化放在全局范围内，避免每个请求都重新加载模型
device = "cuda"
try:
    # 将device和 CartoonGenerator 初始化放在全局范围内，避免每个请求都重新加载模型
    cartoon_generator = CartoonGenerator(device=device)
    # avatar_generator.unet.eval()
    # avatar_generator.vae.eval()
except Exception as e:
    logger.error(f"Failed to initialize avatar generator: {str(e)}")
    raise  # 重新抛出异常，以便在应用启动时处理错误

def bytes_to_pil(bytes_data):
    image = Image.open(BytesIO(bytes_data))
    return image

def bytes_to_image(image: bytes) -> Image.Image:
    image = Image.open(BytesIO(image))
    return image

def swap_face(data):
    image = data.get('data', None)
    style = data.get('style', None)

    image = base64.b64decode(image)
    image = bytes_to_pil(image)

    # 确保所有图像是 RGB 或 L 模式
    final_image = generate_image(image, cartoon_generator, style)
    
    return final_image

def image_to_base64(image: Image.Image) -> str:
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')


# Flask 路由和视图函数
@app.route('/StarDrawingFilter', methods=['POST'])
def StarDrawingFilter():
    data = request.get_json()
    if not data or 'data' not in data or 'style' not in data:
        response_data = {
            "data": None,
            "code": 400
        }
        return jsonify(response_data)
    
    try:
        start = time.time()        
        new_image = swap_face(data)
        processing_time = time.time() - start
        logger.info(f"processing image: {processing_time}")  
        base64_image = image_to_base64(new_image)
        response_data = {
            "data": base64_image,
            "code": 200
        }
        
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}", exc_info=True)
        response_data = {
            "data": None,
            "code": 1001
        }
    
    return jsonify(response_data)


if __name__ == '__main__':
    logger.info("Starting Flask application...")
    app.run(host="0.0.0.0", port=7861, threaded=False)
